# Copyright 2025 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

# Code generated by the Google Gen AI SDK generator DO NOT EDIT.

import datetime
import json
import logging
import time
from typing import Any, Optional, Union
from urllib.parse import urlencode

from google.genai import _api_module
from google.genai import _common
from google.genai import types as genai_types
from google.genai._common import get_value_by_path as getv
from google.genai._common import set_value_by_path as setv

from . import _prompt_optimizer_utils
from . import types


logger = logging.getLogger("vertexai_genai.promptoptimizer")


def _CustomJobParameters_to_vertex(
    from_object: Union[dict[str, Any], object],
    parent_object: Optional[dict[str, Any]] = None,
) -> dict[str, Any]:
    to_object: dict[str, Any] = {}
    if getv(from_object, ["custom_job"]) is not None:
        setv(
            parent_object,
            ["customJob"],
            _CustomJob_to_vertex(getv(from_object, ["custom_job"]), to_object),
        )

    if getv(from_object, ["config"]) is not None:
        setv(to_object, ["config"], getv(from_object, ["config"]))

    return to_object


def _CustomJob_from_vertex(
    from_object: Union[dict[str, Any], object],
    parent_object: Optional[dict[str, Any]] = None,
) -> dict[str, Any]:
    to_object: dict[str, Any] = {}
    if getv(parent_object, ["displayName"]) is not None:
        setv(to_object, ["display_name"], getv(parent_object, ["displayName"]))

    if getv(parent_object, ["jobSpec"]) is not None:
        setv(to_object, ["job_spec"], getv(parent_object, ["jobSpec"]))

    if getv(parent_object, ["encryptionSpec"]) is not None:
        setv(to_object, ["encryption_spec"], getv(parent_object, ["encryptionSpec"]))

    if getv(from_object, ["state"]) is not None:
        setv(to_object, ["state"], getv(from_object, ["state"]))

    if getv(parent_object, ["error"]) is not None:
        setv(to_object, ["error"], getv(parent_object, ["error"]))

    if getv(from_object, ["createTime"]) is not None:
        setv(to_object, ["create_time"], getv(from_object, ["createTime"]))

    if getv(from_object, ["endTime"]) is not None:
        setv(to_object, ["end_time"], getv(from_object, ["endTime"]))

    if getv(from_object, ["labels"]) is not None:
        setv(to_object, ["labels"], getv(from_object, ["labels"]))

    if getv(from_object, ["name"]) is not None:
        setv(to_object, ["name"], getv(from_object, ["name"]))

    if getv(from_object, ["satisfiesPzi"]) is not None:
        setv(to_object, ["satisfies_pzi"], getv(from_object, ["satisfiesPzi"]))

    if getv(from_object, ["satisfiesPzs"]) is not None:
        setv(to_object, ["satisfies_pzs"], getv(from_object, ["satisfiesPzs"]))

    if getv(from_object, ["startTime"]) is not None:
        setv(to_object, ["start_time"], getv(from_object, ["startTime"]))

    if getv(from_object, ["updateTime"]) is not None:
        setv(to_object, ["update_time"], getv(from_object, ["updateTime"]))

    if getv(from_object, ["webAccessUris"]) is not None:
        setv(to_object, ["web_access_uris"], getv(from_object, ["webAccessUris"]))

    return to_object


def _CustomJob_to_vertex(
    from_object: Union[dict[str, Any], object],
    parent_object: Optional[dict[str, Any]] = None,
) -> dict[str, Any]:
    to_object: dict[str, Any] = {}
    if getv(from_object, ["display_name"]) is not None:
        setv(parent_object, ["displayName"], getv(from_object, ["display_name"]))

    if getv(from_object, ["job_spec"]) is not None:
        setv(parent_object, ["jobSpec"], getv(from_object, ["job_spec"]))

    if getv(from_object, ["encryption_spec"]) is not None:
        setv(parent_object, ["encryptionSpec"], getv(from_object, ["encryption_spec"]))

    if getv(from_object, ["state"]) is not None:
        setv(to_object, ["state"], getv(from_object, ["state"]))

    if getv(from_object, ["error"]) is not None:
        setv(parent_object, ["error"], getv(from_object, ["error"]))

    if getv(from_object, ["create_time"]) is not None:
        setv(to_object, ["createTime"], getv(from_object, ["create_time"]))

    if getv(from_object, ["end_time"]) is not None:
        setv(to_object, ["endTime"], getv(from_object, ["end_time"]))

    if getv(from_object, ["labels"]) is not None:
        setv(to_object, ["labels"], getv(from_object, ["labels"]))

    if getv(from_object, ["name"]) is not None:
        setv(to_object, ["name"], getv(from_object, ["name"]))

    if getv(from_object, ["satisfies_pzi"]) is not None:
        setv(to_object, ["satisfiesPzi"], getv(from_object, ["satisfies_pzi"]))

    if getv(from_object, ["satisfies_pzs"]) is not None:
        setv(to_object, ["satisfiesPzs"], getv(from_object, ["satisfies_pzs"]))

    if getv(from_object, ["start_time"]) is not None:
        setv(to_object, ["startTime"], getv(from_object, ["start_time"]))

    if getv(from_object, ["update_time"]) is not None:
        setv(to_object, ["updateTime"], getv(from_object, ["update_time"]))

    if getv(from_object, ["web_access_uris"]) is not None:
        setv(to_object, ["webAccessUris"], getv(from_object, ["web_access_uris"]))

    return to_object


def _GetCustomJobParameters_to_vertex(
    from_object: Union[dict[str, Any], object],
    parent_object: Optional[dict[str, Any]] = None,
) -> dict[str, Any]:
    to_object: dict[str, Any] = {}
    if getv(from_object, ["name"]) is not None:
        setv(to_object, ["_url", "name"], getv(from_object, ["name"]))

    if getv(from_object, ["config"]) is not None:
        setv(to_object, ["config"], getv(from_object, ["config"]))

    return to_object


def _OptimizeConfig_to_vertex(
    from_object: Union[dict[str, Any], object],
    parent_object: Optional[dict[str, Any]] = None,
) -> dict[str, Any]:
    to_object: dict[str, Any] = {}

    if getv(from_object, ["optimization_target"]) is not None:
        setv(
            parent_object,
            ["optimizationTarget"],
            getv(from_object, ["optimization_target"]),
        )

    return to_object


def _OptimizeRequestParameters_to_vertex(
    from_object: Union[dict[str, Any], object],
    parent_object: Optional[dict[str, Any]] = None,
) -> dict[str, Any]:
    to_object: dict[str, Any] = {}
    if getv(from_object, ["content"]) is not None:
        setv(to_object, ["content"], getv(from_object, ["content"]))

    if getv(from_object, ["config"]) is not None:
        setv(
            to_object,
            ["config"],
            _OptimizeConfig_to_vertex(getv(from_object, ["config"]), to_object),
        )

    return to_object


class PromptOptimizer(_api_module.BaseModule):
    """Prompt Optimizer"""

    def _optimize_prompt(
        self,
        *,
        content: Optional[genai_types.ContentOrDict] = None,
        config: Optional[types.OptimizeConfigOrDict] = None,
    ) -> types.OptimizeResponseEndpoint:
        """
        Optimize a single prompt.
        """

        parameter_model = types._OptimizeRequestParameters(
            content=content,
            config=config,
        )

        request_url_dict: Optional[dict[str, str]]
        if not self._api_client.vertexai:
            raise ValueError("This method is only supported in the Vertex AI client.")
        else:
            request_dict = _OptimizeRequestParameters_to_vertex(parameter_model)
            request_url_dict = request_dict.get("_url")
            if request_url_dict:
                path = "tuningJobs:optimizePrompt".format_map(request_url_dict)
            else:
                path = "tuningJobs:optimizePrompt"

        query_params = request_dict.get("_query")
        if query_params:
            path = f"{path}?{urlencode(query_params)}"
        # TODO: remove the hack that pops config.
        request_dict.pop("config", None)

        http_options: Optional[types.HttpOptions] = None
        if (
            parameter_model.config is not None
            and parameter_model.config.http_options is not None
        ):
            http_options = parameter_model.config.http_options

        request_dict = _common.convert_to_dict(request_dict)
        request_dict = _common.encode_unserializable_types(request_dict)

        response = self._api_client.request("post", path, request_dict, http_options)

        response_dict = {} if not response.body else json.loads(response.body)

        return_value = types.OptimizeResponseEndpoint._from_response(
            response=response_dict, kwargs=parameter_model.model_dump()
        )

        self._api_client._verify_response(return_value)
        return return_value

    def _create_custom_job_resource(
        self,
        *,
        custom_job: types.CustomJobOrDict,
        config: Optional[types.VertexBaseConfigOrDict] = None,
    ) -> types.CustomJob:
        """
        Creates a custom job.
        """

        parameter_model = types._CustomJobParameters(
            custom_job=custom_job,
            config=config,
        )

        request_url_dict: Optional[dict[str, str]]
        if not self._api_client.vertexai:
            raise ValueError("This method is only supported in the Vertex AI client.")
        else:
            request_dict = _CustomJobParameters_to_vertex(parameter_model)
            request_url_dict = request_dict.get("_url")
            if request_url_dict:
                path = "customJobs".format_map(request_url_dict)
            else:
                path = "customJobs"

        query_params = request_dict.get("_query")
        if query_params:
            path = f"{path}?{urlencode(query_params)}"
        # TODO: remove the hack that pops config.
        request_dict.pop("config", None)

        http_options: Optional[types.HttpOptions] = None
        if (
            parameter_model.config is not None
            and parameter_model.config.http_options is not None
        ):
            http_options = parameter_model.config.http_options

        request_dict = _common.convert_to_dict(request_dict)
        request_dict = _common.encode_unserializable_types(request_dict)

        response = self._api_client.request("post", path, request_dict, http_options)

        response_dict = {} if not response.body else json.loads(response.body)

        if self._api_client.vertexai:
            response_dict = _CustomJob_from_vertex(response_dict)

        return_value = types.CustomJob._from_response(
            response=response_dict, kwargs=parameter_model.model_dump()
        )

        self._api_client._verify_response(return_value)
        return return_value

    def _get_custom_job(
        self, *, name: str, config: Optional[types.VertexBaseConfigOrDict] = None
    ) -> types.CustomJob:
        """
        Gets a custom job.
        """

        parameter_model = types._GetCustomJobParameters(
            name=name,
            config=config,
        )

        request_url_dict: Optional[dict[str, str]]
        if not self._api_client.vertexai:
            raise ValueError("This method is only supported in the Vertex AI client.")
        else:
            request_dict = _GetCustomJobParameters_to_vertex(parameter_model)
            request_url_dict = request_dict.get("_url")
            if request_url_dict:
                path = "customJobs/{name}".format_map(request_url_dict)
            else:
                path = "customJobs/{name}"

        query_params = request_dict.get("_query")
        if query_params:
            path = f"{path}?{urlencode(query_params)}"
        # TODO: remove the hack that pops config.
        request_dict.pop("config", None)

        http_options: Optional[types.HttpOptions] = None
        if (
            parameter_model.config is not None
            and parameter_model.config.http_options is not None
        ):
            http_options = parameter_model.config.http_options

        request_dict = _common.convert_to_dict(request_dict)
        request_dict = _common.encode_unserializable_types(request_dict)

        response = self._api_client.request("get", path, request_dict, http_options)

        response_dict = {} if not response.body else json.loads(response.body)

        if self._api_client.vertexai:
            response_dict = _CustomJob_from_vertex(response_dict)

        return_value = types.CustomJob._from_response(
            response=response_dict, kwargs=parameter_model.model_dump()
        )

        self._api_client._verify_response(return_value)
        return return_value

    """Prompt Optimizer PO-Data."""

    def _wait_for_completion(self, job_name: str) -> types.CustomJob:

        JOB_COMPLETE_STATES = [
            genai_types.JobState.JOB_STATE_SUCCEEDED,
            genai_types.JobState.JOB_STATE_FAILED,
            genai_types.JobState.JOB_STATE_CANCELLED,
            genai_types.JobState.JOB_STATE_PAUSED,
        ]
        JOB_ERROR_STATES = [
            genai_types.JobState.JOB_STATE_FAILED,
            genai_types.JobState.JOB_STATE_CANCELLED,
        ]

        log_wait = 5
        wait_multiplier = 2
        max_wait_time = 60
        previous_time = time.time()

        job = self._get_custom_job(name=job_name)

        while job.state not in JOB_COMPLETE_STATES:
            current_time = time.time()
            if current_time - previous_time >= log_wait:
                logger.info(f"Waiting for job to complete. Current state: {job.state}")
                log_wait = min(log_wait * wait_multiplier, max_wait_time)
                previous_time = current_time
            time.sleep(log_wait)
            job = self._get_custom_job(name=job_name)

        logger.info(f"Job state: {job.state}")

        if job.state in JOB_ERROR_STATES:
            raise RuntimeError(f"Job failed with state: {job.state}")
        else:
            logger.info(f"Job completed with state: {job.state}")
        return job

    def optimize(
        self,
        method: types.PromptOptimizerMethod,
        config: types.PromptOptimizerConfigOrDict,
    ) -> types.CustomJob:
        """Call PO-Data optimizer.

        Args:
          method: The method for optimizing multiple prompts. Supported methods:
            VAPO, OPTIMIZATION_TARGET_GEMINI_NANO.
          config: PromptOptimizerConfig instance containing the
              configuration for prompt optimization.
        Returns:
          The custom job that was created.
        """

        if isinstance(config, dict):
            config = types.PromptOptimizerConfig(**config)

        if not config.config_path:
            raise ValueError("Config path is required.")

        _OPTIMIZER_METHOD_TO_CONTAINER_URI = {
            types.PromptOptimizerMethod.VAPO: "us-docker.pkg.dev/vertex-ai/cair/vaipo:preview_v1_0",
            types.PromptOptimizerMethod.OPTIMIZATION_TARGET_GEMINI_NANO: "us-docker.pkg.dev/vertex-ai/cair/android-apo:preview_v1_0",
        }
        container_uri = _OPTIMIZER_METHOD_TO_CONTAINER_URI.get(method)
        if not container_uri:
            raise ValueError(
                'Only "VAPO" and "OPTIMIZATION_TARGET_GEMINI_NANO" '
                "methods are currently supported."
            )

        if config.optimizer_job_display_name:
            display_name = config.optimizer_job_display_name
        else:
            timestamp = datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
            display_name = f"{method.value.lower()}-optimizer-{timestamp}"

        wait_for_completion = config.wait_for_completion
        bucket = "/".join(config.config_path.split("/")[:-1])

        region = self._api_client.location
        project = self._api_client.project
        container_args = {
            "config": config.config_path,
        }
        args = ["--%s=%s" % (k, v) for k, v in container_args.items()]
        worker_pool_specs = [
            types.WorkerPoolSpec(
                replica_count=1,
                machine_spec=types.MachineSpec(machine_type="n1-standard-4"),
                container_spec=types.ContainerSpec(
                    image_uri=container_uri,
                    args=args,
                ),
            )
        ]

        service_account = _prompt_optimizer_utils._get_service_account(config)

        job_spec = types.CustomJobSpec(
            worker_pool_specs=worker_pool_specs,
            base_output_directory=genai_types.GcsDestination(output_uri_prefix=bucket),
            service_account=service_account,
        )

        custom_job = types.CustomJob(
            display_name=display_name,
            job_spec=job_spec,
        )

        job = self._create_custom_job_resource(
            custom_job=custom_job,
        )

        # Get the job resource name
        job_resource_name = job.name
        if not job_resource_name:
            raise ValueError(f"Error creating job: {job}")
        job_id = job_resource_name.split("/")[-1]
        logger.info("Job created: %s", job.name)

        # Construct the dashboard URL
        dashboard_url = f"https://console.cloud.google.com/vertex-ai/locations/{region}/training/{job_id}/cpu?project={project}"
        logger.info("View the job status at: %s", dashboard_url)

        if wait_for_completion:
            job = self._wait_for_completion(job_id)
        return job

    def optimize_prompt(
        self,
        *,
        prompt: str,
        config: Optional[types.OptimizeConfig] = None,
    ) -> types.OptimizeResponse:
        """Makes an API request to _optimize_prompt and returns the parsed response.

        Example usage:
        client = vertexai.Client(project=PROJECT_NAME, location='us-central1')
        prompt = "Generate system instructions for analyzing medical articles"
        response = client.prompt_optimizer.optimize_prompt(prompt=prompt)
        print(response.suggested_prompt)

        Args:
          prompt: The prompt to optimize.
          config: Optional.The configuration for prompt optimization. To optimize
            prompts from Android API provide
            types.OptimizeConfig(
                optimization_target=types.OptimizeTarget.OPTIMIZATION_TARGET_GEMINI_NANO
            )
        Returns:
          The parsed response from the API request.
        """

        prompt = genai_types.Content(parts=[genai_types.Part(text=prompt)], role="user")
        # TODO: b/435653980 - replace the custom method with a generated method.
        return self._custom_optimize_prompt(
            content=prompt,
            config=config,
        )

    def _custom_optimize_prompt(
        self,
        *,
        content: Optional[genai_types.ContentOrDict] = None,
        config: Optional[types.OptimizeConfigOrDict] = None,
    ) -> types.OptimizeResponse:
        """Optimize a single prompt.

        Sends a request to the tuningJobs:optimizePrompt streaming endpoint.
        Then gathers the response, concatenates into one string and returns
        the parsed response.
        """

        parameter_model = types._OptimizeRequestParameters(
            content=content,
            config=config,
        )
        request_url_dict: Optional[dict[str, str]]
        if not self._api_client.vertexai:
            raise ValueError("This method is only supported in the Vertex AI client.")
        else:
            request_dict = _OptimizeRequestParameters_to_vertex(parameter_model)
            request_url_dict = request_dict.get("_url")
            if request_url_dict:
                path = "tuningJobs:optimizePrompt".format_map(request_url_dict)
            else:
                path = "tuningJobs:optimizePrompt"

        query_params = request_dict.get("_query")
        if query_params:
            path = f"{path}?{urlencode(query_params)}"
        # TODO: remove the hack that pops config.
        request_dict.pop("config", None)

        http_options: Optional[genai_types.HttpOptions] = None
        if (
            parameter_model.config is not None
            and parameter_model.config.http_options is not None
        ):
            http_options = parameter_model.config.http_options

        request_dict = _common.convert_to_dict(request_dict)
        request_dict = _common.encode_unserializable_types(request_dict)

        response = self._api_client.request("post", path, request_dict, http_options)

        response_list = "" if not response.body else json.loads(response.body)

        return_value = []

        for response_dict in response_list:
            response_value = types.OptimizeResponseEndpoint._from_response(
                response=response_dict, kwargs=parameter_model.model_dump()
            )
            self._api_client._verify_response(response_value)
            if (
                response_value.content is not None
                and len(response_value.content.parts) > 0
                and response_value.content.parts[0].text is not None
            ):
                return_value.append(response_value.content.parts[0].text)

        output = "".join(return_value)
        final_response = types.OptimizeResponse(raw_text_response=output)
        try:
            final_response.parsed_response = _prompt_optimizer_utils._parse(output)
        except Exception as e:
            logger.warning(
                f"Failed to parse response: {e}. Returning only raw_text_response."
            )
        return final_response


class AsyncPromptOptimizer(_api_module.BaseModule):
    """Prompt Optimizer"""

    async def _optimize_prompt(
        self,
        *,
        content: Optional[genai_types.ContentOrDict] = None,
        config: Optional[types.OptimizeConfigOrDict] = None,
    ) -> types.OptimizeResponseEndpoint:
        """
        Optimize a single prompt.
        """

        parameter_model = types._OptimizeRequestParameters(
            content=content,
            config=config,
        )

        request_url_dict: Optional[dict[str, str]]
        if not self._api_client.vertexai:
            raise ValueError("This method is only supported in the Vertex AI client.")
        else:
            request_dict = _OptimizeRequestParameters_to_vertex(parameter_model)
            request_url_dict = request_dict.get("_url")
            if request_url_dict:
                path = "tuningJobs:optimizePrompt".format_map(request_url_dict)
            else:
                path = "tuningJobs:optimizePrompt"

        query_params = request_dict.get("_query")
        if query_params:
            path = f"{path}?{urlencode(query_params)}"
        # TODO: remove the hack that pops config.
        request_dict.pop("config", None)

        http_options: Optional[types.HttpOptions] = None
        if (
            parameter_model.config is not None
            and parameter_model.config.http_options is not None
        ):
            http_options = parameter_model.config.http_options

        request_dict = _common.convert_to_dict(request_dict)
        request_dict = _common.encode_unserializable_types(request_dict)

        response = await self._api_client.async_request(
            "post", path, request_dict, http_options
        )

        response_dict = {} if not response.body else json.loads(response.body)

        return_value = types.OptimizeResponseEndpoint._from_response(
            response=response_dict, kwargs=parameter_model.model_dump()
        )

        self._api_client._verify_response(return_value)
        return return_value

    async def _create_custom_job_resource(
        self,
        *,
        custom_job: types.CustomJobOrDict,
        config: Optional[types.VertexBaseConfigOrDict] = None,
    ) -> types.CustomJob:
        """
        Creates a custom job.
        """

        parameter_model = types._CustomJobParameters(
            custom_job=custom_job,
            config=config,
        )

        request_url_dict: Optional[dict[str, str]]
        if not self._api_client.vertexai:
            raise ValueError("This method is only supported in the Vertex AI client.")
        else:
            request_dict = _CustomJobParameters_to_vertex(parameter_model)
            request_url_dict = request_dict.get("_url")
            if request_url_dict:
                path = "customJobs".format_map(request_url_dict)
            else:
                path = "customJobs"

        query_params = request_dict.get("_query")
        if query_params:
            path = f"{path}?{urlencode(query_params)}"
        # TODO: remove the hack that pops config.
        request_dict.pop("config", None)

        http_options: Optional[types.HttpOptions] = None
        if (
            parameter_model.config is not None
            and parameter_model.config.http_options is not None
        ):
            http_options = parameter_model.config.http_options

        request_dict = _common.convert_to_dict(request_dict)
        request_dict = _common.encode_unserializable_types(request_dict)

        response = await self._api_client.async_request(
            "post", path, request_dict, http_options
        )

        response_dict = {} if not response.body else json.loads(response.body)

        if self._api_client.vertexai:
            response_dict = _CustomJob_from_vertex(response_dict)

        return_value = types.CustomJob._from_response(
            response=response_dict, kwargs=parameter_model.model_dump()
        )

        self._api_client._verify_response(return_value)
        return return_value

    async def _get_custom_job(
        self, *, name: str, config: Optional[types.VertexBaseConfigOrDict] = None
    ) -> types.CustomJob:
        """
        Gets a custom job.
        """

        parameter_model = types._GetCustomJobParameters(
            name=name,
            config=config,
        )

        request_url_dict: Optional[dict[str, str]]
        if not self._api_client.vertexai:
            raise ValueError("This method is only supported in the Vertex AI client.")
        else:
            request_dict = _GetCustomJobParameters_to_vertex(parameter_model)
            request_url_dict = request_dict.get("_url")
            if request_url_dict:
                path = "customJobs/{name}".format_map(request_url_dict)
            else:
                path = "customJobs/{name}"

        query_params = request_dict.get("_query")
        if query_params:
            path = f"{path}?{urlencode(query_params)}"
        # TODO: remove the hack that pops config.
        request_dict.pop("config", None)

        http_options: Optional[types.HttpOptions] = None
        if (
            parameter_model.config is not None
            and parameter_model.config.http_options is not None
        ):
            http_options = parameter_model.config.http_options

        request_dict = _common.convert_to_dict(request_dict)
        request_dict = _common.encode_unserializable_types(request_dict)

        response = await self._api_client.async_request(
            "get", path, request_dict, http_options
        )

        response_dict = {} if not response.body else json.loads(response.body)

        if self._api_client.vertexai:
            response_dict = _CustomJob_from_vertex(response_dict)

        return_value = types.CustomJob._from_response(
            response=response_dict, kwargs=parameter_model.model_dump()
        )

        self._api_client._verify_response(return_value)
        return return_value

    # Todo: b/428953357 - Add example in the README.
    async def optimize(
        self,
        method: types.PromptOptimizerMethod,
        config: types.PromptOptimizerConfigOrDict,
    ) -> types.CustomJob:
        """Call async Vertex AI Prompt Optimizer (VAPO).


        Note: The `wait_for_completion` parameter in the config will be
        ignored when using the AsyncClient, as it is not supported.

        Example usage:
        client = vertexai.Client(project=PROJECT_NAME, location='us-central1')
        vapo_config = vertexai.types.PromptOptimizerConfig(
            config_path='gs://you-bucket-name/your-config.json',
            service_account=service_account,
        )
        job = await client.aio.prompt_optimizer.optimize(
            method=types.PromptOptimizerMethod.VAPO, config=vapo_config)

        Args:
          method: The method for optimizing multiple prompts. Supported methods:
            VAPO, OPTIMIZATION_TARGET_GEMINI_NANO.
          config: PromptOptimizerConfig instance containing the
            configuration for prompt optimization.
        Returns:
          The custom job that was created.
        """
        if isinstance(config, dict):
            config = types.PromptOptimizerConfig(**config)

        if not config.config_path:
            raise ValueError("Config path is required.")

        _OPTIMIZER_METHOD_TO_CONTAINER_URI = {
            types.PromptOptimizerMethod.VAPO: "us-docker.pkg.dev/vertex-ai/cair/vaipo:preview_v1_0",
            types.PromptOptimizerMethod.OPTIMIZATION_TARGET_GEMINI_NANO: "us-docker.pkg.dev/vertex-ai/cair/android-apo:preview_v1_0",
        }
        container_uri = _OPTIMIZER_METHOD_TO_CONTAINER_URI.get(method)
        if not container_uri:
            raise ValueError(
                'Only "VAPO" and "OPTIMIZATION_TARGET_GEMINI_NANO" '
                "methods are currently supported."
            )

        if config.wait_for_completion:
            logger.info(
                "Ignoring wait_for_completion=True since the AsyncClient does not support it."
            )

        if config.optimizer_job_display_name:
            display_name = config.optimizer_job_display_name
        else:
            timestamp = datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
            display_name = f"{method.value.lower()}-optimizer-{timestamp}"

        if not config.config_path:
            raise ValueError("Config path is required.")
        bucket = "/".join(config.config_path.split("/")[:-1])

        region = self._api_client.location
        project = self._api_client.project
        container_args = {
            "config": config.config_path,
        }
        args = ["--%s=%s" % (k, v) for k, v in container_args.items()]
        worker_pool_specs = [
            types.WorkerPoolSpec(
                replica_count=1,
                machine_spec=types.MachineSpec(machine_type="n1-standard-4"),
                container_spec=types.ContainerSpec(
                    image_uri=container_uri,
                    args=args,
                ),
            )
        ]

        service_account = _prompt_optimizer_utils._get_service_account(config)

        job_spec = types.CustomJobSpec(
            worker_pool_specs=worker_pool_specs,
            base_output_directory=genai_types.GcsDestination(output_uri_prefix=bucket),
            service_account=service_account,
        )

        custom_job = types.CustomJob(
            display_name=display_name,
            job_spec=job_spec,
        )

        job = await self._create_custom_job_resource(
            custom_job=custom_job,
        )

        # Get the job id for the dashboard url and display to the user.
        job_resource_name = job.name
        if not job_resource_name:
            raise ValueError(f"Error creating job: {job}")
        job_id = job_resource_name.split("/")[-1]
        logger.info("Job created: %s", job.name)

        # Construct the dashboard URL to show to the user.
        dashboard_url = f"https://console.cloud.google.com/vertex-ai/locations/{region}/training/{job_id}/cpu?project={project}"
        logger.info("View the job status at: %s", dashboard_url)

        return job

    async def _custom_optimize_prompt(
        self,
        *,
        content: Optional[genai_types.ContentOrDict] = None,
        config: Optional[types.OptimizeConfigOrDict] = None,
    ) -> types.OptimizeResponse:
        """Optimize a single prompt."""

        parameter_model = types._OptimizeRequestParameters(
            content=content,
            config=config,
        )
        request_url_dict: Optional[dict[str, str]]
        if not self._api_client.vertexai:
            raise ValueError("This method is only supported in the Vertex AI client.")
        else:
            request_dict = _OptimizeRequestParameters_to_vertex(parameter_model)
            request_url_dict = request_dict.get("_url")
            if request_url_dict:
                path = "tuningJobs:optimizePrompt".format_map(request_url_dict)
            else:
                path = "tuningJobs:optimizePrompt"

        query_params = request_dict.get("_query")
        if query_params:
            path = f"{path}?{urlencode(query_params)}"
        # TODO: remove the hack that pops config.
        request_dict.pop("config", None)

        http_options: Optional[types.HttpOptions] = None
        if (
            parameter_model.config is not None
            and parameter_model.config.http_options is not None
        ):
            http_options = parameter_model.config.http_options

        request_dict = _common.convert_to_dict(request_dict)
        request_dict = _common.encode_unserializable_types(request_dict)

        response = await self._api_client.async_request(
            "post", path, request_dict, http_options
        )

        response_list = "" if not response.body else json.loads(response.body)

        return_value = []

        for response_dict in response_list:
            response_value = types.OptimizeResponseEndpoint._from_response(
                response=response_dict, kwargs=parameter_model.model_dump()
            )
            self._api_client._verify_response(response_value)
            if (
                response_value.content is not None
                and len(response_value.content.parts) > 0
                and response_value.content.parts[0].text is not None
            ):
                return_value.append(response_value.content.parts[0].text)

        output = "".join(return_value)
        final_response = types.OptimizeResponse(raw_text_response=output)
        try:
            final_response.parsed_response = _prompt_optimizer_utils._parse(output)
        except Exception as e:
            logger.warning(
                f"Failed to parse response: {e}. Returning only raw_text_response."
            )
        return final_response

    async def optimize_prompt(
        self,
        *,
        prompt: str,
        config: Optional[types.OptimizeConfig] = None,
    ) -> types.OptimizeResponse:
        """Makes an async request to _optimize_prompt and returns an optimized prompt.

        Example usage:
        client = vertexai.Client(project=PROJECT_NAME, location='us-central1')
        prompt = "Generate system instructions for analyzing medical articles"
        response = await client.aio.prompt_optimizer.optimize_prompt(prompt=prompt)

        Args:
          prompt: The prompt to optimize.
          config: Optional.The configuration for prompt optimization. To optimize
            prompts from Android API provide
            types.OptimizeConfig(
                optimization_target=types.OptimizeTarget.OPTIMIZATION_TARGET_GEMINI_NANO
            )
        Returns:
          The parsed response from the API request.
        """

        prompt = genai_types.Content(parts=[genai_types.Part(text=prompt)], role="user")
        # TODO: b/435653980 - replace the custom method with a generated method.
        return await self._custom_optimize_prompt(
            content=prompt,
            config=config,
        )
